import torch
import torch.nn as nn
import torch.nn.functional as F


class ConvBlock(nn.Module):
    def __init__(self, in_channels, out_channels, stride=2):
        super(ConvBlock, self).__init__()
        self.conv = nn.Sequential(
            nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=2, stride=stride),
        )

    def forward(self, x):
        return self.conv(x)


class CustomAlexNet(nn.Module):
    def __init__(self):
        super(CustomAlexNet, self).__init__()

        # First input branch for 128x128 images
        self.branch_128 = nn.Sequential(
            ConvBlock(3, 64, stride=2),  # output will be 64x64
            ConvBlock(64, 128, stride=2),  # output will be 32x32
            ConvBlock(128, 128, stride=1),  # output will be 16x16
        )

        # Second input branch for 256x256 images
        self.branch_256 = nn.Sequential(
            ConvBlock(3, 64, stride=2),  # output will be 128x128
            ConvBlock(64, 128, stride=2),  # output will be 64x64
            ConvBlock(128, 256, stride=2),  # output will be 32x32
        )

        # Merge layer
        self.merge_conv = ConvBlock(128 + 256, 256, stride=1)  # after concatenation along the channel dimension

        # Final layers
        self.avgpool = nn.AdaptiveAvgPool2d((1, 1))  # Output will be 1x1
        self.fc = nn.Linear(256, 4096)  # Adjusted to the output size of the adaptive avg pool
        self.out = nn.Linear(4096, 10)  # Output layer for 10 classes

    def forward(self, x_128, x_256):
        # Process the two branches
        x_128 = self.branch_128(x_128)
        x_256 = self.branch_256(x_256)

        # Concatenate along the channel dimension
        x = torch.cat((x_128, x_256), 1)

        # Merge layer
        x = self.merge_conv(x)

        # Final layers
        x = self.avgpool(x)
        x = x.view(x.size(0), -1)  # Flatten the tensor
        x = self.fc(x)
        x = F.relu(x)
        x = self.out(x)

        return x


# Verify the model can run
model = CustomAlexNet()
# Create dummy inputs
dummy_input_128 = torch.randn(1, 3, 128, 128)
dummy_input_256 = torch.randn(1, 3, 256, 256)
# Run the model (forward pass)
output = model(dummy_input_128, dummy_input_256)
print(output)  # Print the output to verify that it runs